#include "config.h"

#include "lstatfit.h"

extern int exp_dist_f(const gsl_vector *x, void *data, gsl_vector *f);
extern int exp_dist_df(const gsl_vector *x, void *data, gsl_matrix *J);
extern void exp_dist_report(FILE *fp, msl_fit_t *fit, ecm_uint32 xdiff);
extern ecm_double exp_goodness(msl_fit_t *fit);

extern int beta_dist_f(const gsl_vector *x, void *data, gsl_vector *f);
extern int beta_dist_df(const gsl_vector *x, void *data, gsl_matrix *J);
extern void beta_dist_report(FILE *fp, msl_fit_t *fit, ecm_uint32 xdiff);
extern ecm_double beta_goodness(msl_fit_t *fit);

extern int lognormal_dist_f(const gsl_vector *x, void *data, gsl_vector *f);
extern int lognormal_dist_df(const gsl_vector *x, void *data, gsl_matrix *J);
extern void lognormal_dist_report(FILE *fp, msl_fit_t *fit, ecm_uint32 xdiff);
extern ecm_double lognormal_goodness(msl_fit_t *fit);

extern void weibull_dist_report(FILE *fp, msl_fit_t *fit, ecm_uint32 xdiff);
extern int weibull_dist_f(const gsl_vector *x, void *data, gsl_vector *f);
extern int weibull_dist_df(const gsl_vector *x, void *data, gsl_matrix *J);
extern ecm_double weibull_goodness(msl_fit_t *fit);

static fit_info_t	fit_infos[] = {
	{
		"exponential",
		exp_dist_f,
		exp_dist_df,
		exp_dist_report,
		exp_goodness,
		1,
		{ 0.5 },
		1e-8
	},
	{
		"beta",
		beta_dist_f,
		beta_dist_df,
		beta_dist_report,
		beta_goodness,
		3,
		{ 40, 0.5, 0.5 },
		1e-6
	},
	{
		"lognormal",
		lognormal_dist_f,
		lognormal_dist_df,
		lognormal_dist_report,
		lognormal_goodness,
		2,
		{ 0.1, 50 },
		1e-4
	},
	{
		"weibull",
		weibull_dist_f,
		weibull_dist_df,
		weibull_dist_report,
		weibull_goodness,
		2,
		{ 2, 1 },
		1e-4
	}
};

static FILE *
setup_output(void)
{
	FILE	*fp;

	fp = ecm_fopen("output.txt", "w");
	if (fp == NULL) {
		return NULL;
	}

	return fp;
}

static int
dist_fdf(const gsl_vector *x, void *data, gsl_vector *f, gsl_matrix *J)
{
	msl_fit_t	*fit = (msl_fit_t *)data;

	fit->info->dist_f(x, data, f);
	fit->info->dist_df(x, data, J);

	return GSL_SUCCESS;
}

msl_fit_t *
lstat_fitting(const char *name, msl_fit_type_t fit_type, ecm_uint32 n_inputs, ecm_double *xvalues, ecm_double *yvalues)
{
	msl_fit_t	*fit;
	fit_info_t	*info = &fit_infos[fit_type - 1];
	gsl_vector_view		x;

	fit = (msl_fit_t *)ecm_malloc(sizeof(msl_fit_t));

	fit->name = ecm_strdup(name);
	fit->info = info;

	fit->fit_type = fit_type;
	fit->covar = gsl_matrix_alloc(info->p, info->p);
	fit->n = n_inputs;
	fit->xvalues = xvalues;
	fit->yvalues = yvalues;

	x = gsl_vector_view_array(info->x_init, info->p);
	fit->f.f = info->dist_f;
	fit->f.df = info->dist_df;
	fit->f.fdf = &dist_fdf;
	fit->f.n = n_inputs;
	fit->f.p = info->p;
	fit->f.params = fit;

	fit->xrange = xvalues[fit->n - 1] - xvalues[0];

	fit->T = gsl_multifit_fdfsolver_lmsder;
	fit->s = gsl_multifit_fdfsolver_alloc(fit->T, fit->n, info->p);
	gsl_multifit_fdfsolver_set(fit->s, &fit->f, &x.vector);

	fit->iter = 0;
	fit->status = 0;

	do {
		fit->iter++;
		fit->status = gsl_multifit_fdfsolver_iterate(fit->s);
     
		if (fit->status) {
			printf("status = %s\n", gsl_strerror(fit->status));
			return NULL;
		}
     
		fit->status = gsl_multifit_test_delta(fit->s->dx, fit->s->x, fit->info->tolerance, fit->info->tolerance);
	} while (fit->status == GSL_CONTINUE && fit->iter < 500);

#if 0 ////TODO
	gsl_multifit_covar(fit->s->J, 0.0, fit->covar);

	{
		double chi = gsl_blas_dnrm2(fit->s->f);
		double dof = fit->n - info->p;

		printf("%lf\n", pow(chi, 2.0) / dof);
	}
#endif

	if (fit->status == 0) {
		printf("iteration: %d\n", fit->iter);
	}

	return fit;
}

void
lstat_fitting_free(msl_fit_t *fit)
{
	if (fit == NULL)
		return;

	ecm_free(fit->name);
	gsl_matrix_free(fit->covar);
	ecm_free(fit);
}

void
lstat_report_fitting(msl_fit_t *fit, ecm_uint32 xdiff)
{
	FILE	*fp;

	if ((fp = setup_output()) == NULL)
		return;

	ecm_fprintf(fp, "# name: %s\n", fit->name);
	ecm_fprintf(fp, "# model: %s\n", fit->info->desc);

	fit->info->dist_report(fp, fit, xdiff);

	ecm_fclose(fp);
}

ecm_double
lstat_get_goodness(msl_fit_t *fit)
{
	ecm_double	sq_diffs;

	sq_diffs = fit->info->dist_goodness(fit);
	return sq_diffs / fit->n;
}

msl_fit_type_t
parse_model_type(const char *typestr)
{
	ecm_uint32	i;

	for (i = 0; i < MSL_FIT_MAX; i++) {
		if (ecm_strcmp(fit_infos[i].desc, typestr) == 0) {
			return (i + 1);
		}
	}
	return MSL_FIT_NONE;
}
